/*************************************************************************
 *
 * This file is part of the SAMRAI distribution.  For full copyright
 * information, see COPYRIGHT and LICENSE.
 *
 * Copyright:     (c) 1997-2021 Lawrence Livermore National Security, LLC
 * Description:   Main program for modified-Bratu problem
 *
 ************************************************************************/

#include "SAMRAI/SAMRAI_config.h"

#include <string>
#include <sys/stat.h>
#include <string>

/*
 * Headers for basic SAMRAI objects used in this sample code.
 */
#include "SAMRAI/tbox/Database.h"
#include "SAMRAI/tbox/BalancedDepthFirstTree.h"
#include "SAMRAI/hier/BaseGridGeometry.h"
#include "SAMRAI/tbox/InputManager.h"
#include "SAMRAI/tbox/PIO.h"
#include "SAMRAI/tbox/RestartManager.h"
#include "SAMRAI/tbox/SAMRAIManager.h"
#include "SAMRAI/tbox/SAMRAI_MPI.h"
#include "SAMRAI/tbox/Utilities.h"
#include "SAMRAI/tbox/Timer.h"
#include "SAMRAI/tbox/TimerManager.h"
#include "SAMRAI/hier/VariableDatabase.h"

/*
 * Headers for major algorithm/data structure objects from SAMRAI
 */
#include "SAMRAI/mesh/BergerRigoutsos.h"
#include "SAMRAI/geom/CartesianGridGeometry.h"
#include "SAMRAI/mesh/GriddingAlgorithm.h"
#include "SAMRAI/mesh/TreeLoadBalancer.h"
#include "SAMRAI/hier/PatchHierarchy.h"
#include "SAMRAI/mesh/StandardTagAndInitialize.h"
#include "SAMRAI/appu/VisItDataWriter.h"

/*
 * Headers for major application-specific and solver testing objects
 */
#include "SAMRAI/algs/ImplicitIntegrator.h"
#ifdef HAVE_PETSC
#include "SAMRAI/solv/SNES_SAMRAIContext.h"
#endif
#ifdef HAVE_SUNDIALS
#include "SAMRAI/solv/KINSOL_SAMRAIContext.h"
#endif
#ifdef HAVE_HYPRE
#include "SAMRAI/solv/CellPoissonHypreSolver.h"
#endif
#include "ModifiedBratuProblem.h"
#include "SAMRAI/solv/NonlinearSolverStrategy.h"


#include <vector>
#include <memory>

using namespace SAMRAI;

/*!
 * @brief Main program for nonlinear example.
 *
 * This is the main program for a sample application that employs
 * implicit time integration on an AMR patch hierarchy and requires
 * a nonlinear solver to advance the solution at each time step.
 * The problem to be solved is:
 *
 *        du/dt = div(D(x)*grad(u)) + lambda*exp(u) + f(x,u)
 *
 * The application program is constructed by composing a variety of
 * algorithm objects found in SAMRAI plus numerical operations that
 * are specific to this application.  The following discussion
 * summarizes these objects.
 *
 *    hier::PatchHierarchy - A container for the AMR patch
 *        hierarchy and the data on the grid.
 *
 *    geom::CartesianGridGeometry - Defines and maintains the
 *       Cartesian coordinate system on the grid.  The
 *       hier::PatchHierarchy maintains a reference to this
 *       object.
 *
 * The implicit integrator algorithm uses these two components to
 * integrate data one timestep.  Time-stepping is done in the main
 * routine. This integrator uses the same timestep across all patches,
 * regardless of the refinement.
 *
 *    algs::ImplicitIntegrator - advances solution on patches.
 *             Controls integration of quantities over entire patch
 *             hierarchy. Takes the user-defined class that defines
 *             patch operations (ModifiedBratuProblem in this case)
 *           as an argument.
 *
 * The user-supplied object defines characteristics and operations to
 * be carried out on each patch.  We pass this object to the integrator
 * which will orchestrate the integration on all patches.
 *
 *    ModifiedBratuProblem - Defines variables and numerical routines
 *            on the hierarchy.  This includes routines needed by the
 *          implicit integrator as well as nonlinear solvers used
 *          by the integrator.
 *
 *    mesh::GriddingAlgorithm - Drives the AMR patch hierarchy
 *       generation and regridding procedures.  This object maintains
 *       references to three other algorithmic objects with
 *       which it is configured when they are passed into its
 *       constructor.   They are:
 *
 *       mesh::BergerRigoutsos - Clusters cells tagged for
 *          refinement on a patch level into a collection of
 *          logically-rectangular box domains.
 *
 *       mesh::LoadBalancer - Processes the boxes generated by
 *          the mesh::BergerRigoutsos algorithm into a
 *          configuration from
 *          which patches are contructed.  The routines in this
 *          class assume a spatially-uniform workload distribution;
 *          thus, they attempt to produce a collection of boxes
 *          each of which contains the same number of cells.  The
 *          load balancer also assigns patches to processors.
 *
 *       mesh::StandardTagAndInitialize -
 *          Couples the gridding algorithm to
 *          the ModifiedBratuProblem class which selects cells for
 *          refinement based on criteria defined specifically for
 *          the problem.  This object maintains a pointer to the
 *          ModifiedBratuProblem object, which is passed into its
 *          constructor, for this purpose.
 *
 */

int main(
   int argc,
   char* argv[])
{

   /*
    * Initialize MPI, SAMRAI.
    */

   tbox::SAMRAI_MPI::init(&argc, &argv);
#ifdef HAVE_PETSC
   PetscInitialize(&argc, &argv, PETSC_NULL, PETSC_NULL);
#endif
   tbox::SAMRAIManager::initialize();
   tbox::SAMRAIManager::startup();

   /*
    * Create block to force pointer deallocation.  If this is not done
    * then there will be memory leaks reported.
    */
   {

#if !defined(HAVE_PETSC) || !defined(HAVE_SUNDIALS) || !defined(HAVE_HYPRE)
      tbox::pout << "This example requires the packages PETSC, SUNDIALS, "
                 << "\nand HYPRE to work properly.  SAMRAI was not configured"
                 << "\nwith one or more of these packages."
                 << std::endl;
#else

      /*
       * Process command line arguments.  For each run, the input
       * filename must be specified.  Usage is:
       *
       *     executable <input file name>
       */
      std::string input_filename;

      if (argc != 3) {
         TBOX_ERROR("USAGE:  " << argv[0] << " <input file> -skip_petscrc \n"
                               << "  options:\n"
                               << "  none at this time" << std::endl);
      } else {
         input_filename = argv[1];
      }

      /*
       * Create input database and parse all data in input file.
       */

      std::shared_ptr<tbox::InputDatabase> input_db(
         new tbox::InputDatabase("input_db"));
      tbox::InputManager::getManager()->parseInputFile(input_filename, input_db);

      /*
       * Retrieve "Main" section of the input database.  Read
       * dump information, which is used for writing plot files.
       */

      std::shared_ptr<tbox::Database> main_db(input_db->getDatabase("Main"));

      const tbox::Dimension dim(static_cast<unsigned short>(main_db->getInteger("dim")));

      std::string base_name = "default";
      base_name = main_db->getStringWithDefault("base_name", base_name);

      std::string log_file_name = base_name + ".log";
      tbox::PIO::logOnlyNodeZero(log_file_name);
      tbox::plog << "input_filename = " << input_filename << std::endl;

      int viz_dump_interval = 0;
      if (main_db->keyExists("viz_dump_interval")) {
         viz_dump_interval = main_db->getInteger("viz_dump_interval");
      }

      std::string visit_dump_dirname;
      bool uses_visit = false;
      int visit_number_procs_per_file = 1;
      if (viz_dump_interval > 0) {
         uses_visit = true;
         std::string viz_dump_dirname;
         if (main_db->keyExists("viz_dump_dirname")) {
            viz_dump_dirname = main_db->getString("viz_dump_dirname");
         }
         visit_dump_dirname = viz_dump_dirname;
         if (viz_dump_dirname.empty()) {
            TBOX_ERROR("main(): "
               << "\nviz_dump_dirname is null ... "
               << "\nThis must be specified for use with VisIt"
               << std::endl);
         }
         if (main_db->keyExists("visit_number_procs_per_file")) {
            visit_number_procs_per_file =
               main_db->getInteger("visit_number_procs_per_file");
         }
      }

      int regrid_interval = 0;
      if (main_db->keyExists("regrid_interval")) {
         regrid_interval = main_db->getInteger("regrid_interval");
      }

      std::string nonlinear_solver_package = "KINSOL";
      if (main_db->keyExists("nonlinear_solver_package")) {
         nonlinear_solver_package =
            main_db->getString("nonlinear_solver_package");
      }

      /*
       * Setup the timer manager to trace timing statistics during execution
       * of the code.  The list of timers is given in the tbox::TimerManager
       * section of the input file.
       */
      tbox::TimerManager::createManager(input_db->getDatabase("TimerManager"));
      std::shared_ptr<tbox::Timer> t_all =
         tbox::TimerManager::getManager()->getTimer("appu::main::all");
      t_all->start();

      /*
       * Create major algorithm and data objects which comprise application.
       * Each object is initialized from input data.
       * Refer to each class constructor
       * for details.  For more information on the composition of objects
       * and the roles they play in this application, see comments at top
       * of file.
       */

      std::shared_ptr<appu::VisItDataWriter> visit_data_writer;
      if (uses_visit) {
         visit_data_writer.reset(
            new appu::VisItDataWriter(dim,
               "Bratu VisIt Writer",
               visit_dump_dirname,
               visit_number_procs_per_file));
      }

      std::shared_ptr<geom::CartesianGridGeometry> grid_geometry(
         new geom::CartesianGridGeometry(
            dim,
            "CartesianGeometry",
            input_db->getDatabase("CartesianGeometry")));

      std::shared_ptr<hier::PatchHierarchy> patch_hierarchy(
         new hier::PatchHierarchy(
            "PatchHierarchy",
            grid_geometry,
            input_db->getDatabase("PatchHierarchy")));

      std::string mod_bratu_prob_name = "ModifiedBratuProblem";
      std::string fac_solver_name = mod_bratu_prob_name + ":FAC_solver";
      std::string fac_ops_name = fac_solver_name + "::fac_ops";
      std::string fac_precond_name = fac_solver_name + "::fac_precond";
      std::string hypre_poisson_name = fac_ops_name + "::hypre_solver";

#ifdef HAVE_HYPRE
      std::shared_ptr<solv::CellPoissonHypreSolver> hypre_poisson(
         new solv::CellPoissonHypreSolver(
            dim,
            hypre_poisson_name,
            input_db->isDatabase("hypre_solver") ?
            input_db->getDatabase("hypre_solver") :
            std::shared_ptr<tbox::Database>()));

      std::shared_ptr<solv::CellPoissonFACOps> fac_ops(
         new solv::CellPoissonFACOps(
            hypre_poisson,
            dim,
            fac_ops_name,
            input_db->isDatabase("fac_ops") ?
            input_db->getDatabase("fac_ops") :
            std::shared_ptr<tbox::Database>()));
#else
      std::shared_ptr<solv::CellPoissonFACOps> fac_ops(
         new solv::CellPoissonFACOps(
            dim,
            fac_ops_name,
            input_db->isDatabase("fac_ops") ?
            input_db->getDatabase("fac_ops") :
            std::shared_ptr<tbox::Database>()));
#endif

      std::shared_ptr<solv::FACPreconditioner> fac_precond(
         new solv::FACPreconditioner(
            fac_precond_name,
            fac_ops,
            input_db->isDatabase("fac_precond") ?
            input_db->getDatabase("fac_precond") :
            std::shared_ptr<tbox::Database>()));

      std::shared_ptr<solv::CellPoissonFACSolver> fac_solver(
         new solv::CellPoissonFACSolver(
            dim,
            fac_solver_name,
            fac_precond,
            fac_ops,
            input_db->isDatabase("fac_solver") ?
            input_db->getDatabase("fac_solver") :
            std::shared_ptr<tbox::Database>()));

      ModifiedBratuProblem* bratu_model = new ModifiedBratuProblem(
            mod_bratu_prob_name,
            dim,
            fac_solver,
            input_db->getDatabase("ModifiedBratuProblem"),
            grid_geometry,
            visit_data_writer);

      solv::NonlinearSolverStrategy* nonlinear_solver = 0;

      if (nonlinear_solver_package == "PETSc-SNES") {

#ifdef HAVE_PETSC
         nonlinear_solver =
            new solv::SNES_SAMRAIContext("SNESSolver",
               bratu_model,
               input_db->getDatabase("SNESSolver"));
#else
         TBOX_ERROR("Cannot use PETSc-SNES option because SAMRAI was\n"
            << "not configured to use it.");
#endif

      } else if (nonlinear_solver_package == "KINSOL") {

#ifdef HAVE_SUNDIALS
         nonlinear_solver =
            new solv::KINSOL_SAMRAIContext("KINSOLSolver",
               bratu_model,
               input_db->getDatabase("KINSOLSolver"));
#else
         TBOX_ERROR("Cannot use KINSOL option because SAMRAI was\n"
            << "not configured to use it.");
#endif

      } else {

         TBOX_ERROR("Input key `nonlinear_solver_package' == "
            << nonlinear_solver_package
            << " found in input file is not recognized.");
      }

      algs::ImplicitIntegrator* imp_integrator = new algs::ImplicitIntegrator(
            "ImplicitIntegrator",
            input_db->getDatabase("ImplicitIntegrator"),
            bratu_model,
            nonlinear_solver,
            patch_hierarchy);

      std::shared_ptr<mesh::StandardTagAndInitialize> error_detector(
         new mesh::StandardTagAndInitialize(
            "CellTaggingMethod",
            bratu_model,
            input_db->getDatabase("StandardTagAndInitialize")));

      std::shared_ptr<mesh::BergerRigoutsos> box_generator(
         new mesh::BergerRigoutsos(dim,
            input_db->getDatabase("BergerRigoutsos")));

      std::shared_ptr<mesh::TreeLoadBalancer> load_balancer(
         new mesh::TreeLoadBalancer(
            dim,
            "LoadBalancer",
            input_db->getDatabase("LoadBalancer")));
      load_balancer->setSAMRAI_MPI(tbox::SAMRAI_MPI::getSAMRAIWorld());

      std::shared_ptr<mesh::GriddingAlgorithm> gridding_algorithm(
         new mesh::GriddingAlgorithm(
            patch_hierarchy,
            "GriddingAlgorithm",
            input_db->getDatabase("GriddingAlgorithm"),
            error_detector,
            box_generator,
            load_balancer));

      /*
       * Tag buffers are passed to the gridding algorithm for buffering
       * tagged cells before new levels are created.
       */
      std::vector<int> tag_buffer(patch_hierarchy->getMaxNumberOfLevels());
      for (int ln = 0; ln < patch_hierarchy->getMaxNumberOfLevels(); ++ln) {
         tag_buffer[ln] = 0;
      }
      if (main_db->keyExists("tag_buffer")) {
         std::vector<int> input_tags = main_db->getIntegerVector("tag_buffer");
         if (input_tags.size() > 0) {
            for (int ln0 = 0; ln0 < patch_hierarchy->getMaxNumberOfLevels();
                 ++ln0) {
               if (static_cast<int>(input_tags.size()) > ln0) {
                  tag_buffer[ln0] = ((input_tags[ln0] > 0) ?
                                     input_tags[ln0] : 0);
               } else {
                  int input_tags_size = static_cast<int>(input_tags.size());
                  tag_buffer[ln0] = ((input_tags[input_tags_size - 1] > 0)
                                     ? input_tags[input_tags_size - 1] : 0);
               }
            }
         }
      }

      /*
       * After creating all objects and initializing their state, we
       * print the input database and variable database contents to
       * the log file.
       */

      tbox::plog << "\nCheck input data and variables before simulation:"
                 << std::endl;
      tbox::plog << "Input database..." << std::endl;
      input_db->printClassData(tbox::plog);
      tbox::plog << "\nVariable database..." << std::endl;
      hier::VariableDatabase::getDatabase()->printClassData(tbox::plog);

      /*
       * Initialize AMR hierarchy configuration and data on all patches
       * at initial time.
       * Also, write initial state for
       * vizualization.  Note that we also start a timer for the
       * simulation part of the main program.
       */

      std::shared_ptr<tbox::Timer> main_timer(
         tbox::TimerManager::getManager()->getTimer("apps::main::main"));
      std::shared_ptr<tbox::Timer> solve_timer(
         tbox::TimerManager::getManager()->getTimer("apps::main::solve"));

      main_timer->start();

      double sim_time = imp_integrator->getInitialTime();

      if (tbox::RestartManager::getManager()->isFromRestart()) {

         patch_hierarchy->initializeHierarchy();

         gridding_algorithm->getTagAndInitializeStrategy()->
         resetHierarchyConfiguration(patch_hierarchy,
            0,
            patch_hierarchy->getFinestLevelNumber());

      } else {

         gridding_algorithm->makeCoarsestLevel(sim_time);

         bool done = false;
         bool initial_cycle = true;
         for (int lnum = 0;
              patch_hierarchy->levelCanBeRefined(lnum) && !done; ++lnum) {
            gridding_algorithm->makeFinerLevel(
               tag_buffer[lnum],
               initial_cycle,
               imp_integrator->getIntegratorStep(),
               sim_time);
            done = !(patch_hierarchy->finerLevelExists(lnum));
         }

      }

      /*
       * Now that hierarchy is constructed, initialize vector weights.
       */

      bratu_model->setVectorWeights(patch_hierarchy);

      tbox::RestartManager::getManager()->closeRestartFile();

      if (uses_visit) {
         visit_data_writer->writePlotData(patch_hierarchy,
            imp_integrator->getIntegratorStep(),
            sim_time);
      }

      /*
       * Initialize implicit integrator.  Then, loop over a sequence of time
       * steps until the final simulation time is reached, or the maximum
       * number of steps is exceeded.
       */

      imp_integrator->initialize();

      double dt = imp_integrator->getCurrentDt();
      double end_time = imp_integrator->getFinalTime();

      bool first_step = true;

      while ((sim_time < end_time) &&
             imp_integrator->stepsRemaining()) {

         int iteration_num = imp_integrator->getIntegratorStep() + 1;
         tbox::pout << "++++++++++++++++++++++++++++++++++++++++++++" << std::endl;
         tbox::pout << "At begining of timestep # " << iteration_num - 1
                    << std::endl;
         tbox::pout << "Simulation time is " << sim_time << std::endl;
         tbox::pout << "Time increment is " << dt << std::endl;

         solve_timer->start();
         int solver_retcode = imp_integrator->advanceSolution(dt, first_step);
         solve_timer->stop();

         solv::SNES_SAMRAIContext* snes_solver =
            dynamic_cast<solv::SNES_SAMRAIContext *>(nonlinear_solver);
         if (snes_solver) {
            int nonlinear_itns = snes_solver->getNumberOfNonlinearIterations();
            int linear_itns = snes_solver->getTotalNumberOfLinearIterations();
            tbox::plog << " Nonlinear iterations:  " << nonlinear_itns
                       << " Linear iterations:     " << linear_itns << std::endl;
         }

         bool good_solution = imp_integrator->checkNewSolution(solver_retcode);

         if (good_solution) {
            sim_time = imp_integrator->updateSolution();

            tbox::pout << "At end of timestep # " << iteration_num - 1 << std::endl;
            tbox::pout << "Simulation time is " << sim_time << std::endl;
            tbox::pout << "++++++++++++++++++++++++++++++++++++++++++++"
                       << std::endl;

            /*
             * If desired, write plot file.
             */

            if ((viz_dump_interval > 0) &&
                ((iteration_num % viz_dump_interval) == 0)) {
               if (uses_visit) {
                  visit_data_writer->writePlotData(patch_hierarchy,
                     iteration_num,
                     sim_time);
               }
            }

            /*
             * If desired, regrid patch hierarchy and reset vector weights.
             */

            if ((regrid_interval > 0)
                && ((iteration_num % regrid_interval) == 0)) {
               gridding_algorithm->regridAllFinerLevels(0,
                  tag_buffer,
                  iteration_num,
                  sim_time);
#if defined(HAVE_PETSC) || defined(HAVE_SUNDIALS)
               bratu_model->setVectorWeights(patch_hierarchy);
#endif
               first_step = true;
            } else {
               first_step = false;
            }
         } else {

            tbox::pout << "At end of timestep # " << iteration_num - 1 << std::endl;
            tbox::pout << "Failed to advance solution to " << sim_time << std::endl;
            tbox::pout << "++++++++++++++++++++++++++++++++++++++++++++"
                       << std::endl;
            break;

         }

         dt = imp_integrator->getNextDt(good_solution, solver_retcode);

      }

      /*
       * At conclusion of simulation, stop timer and deallocate objects.
       */
      main_timer->stop();

      tbox::TimerManager::getManager()->print(tbox::plog);

      input_db.reset();
      patch_hierarchy.reset();
      grid_geometry.reset();

      box_generator.reset();
      error_detector.reset();
      load_balancer.reset();
      gridding_algorithm.reset();

      if (imp_integrator)
         delete imp_integrator;
      if (nonlinear_solver)
         delete nonlinear_solver;
      if (bratu_model)
         delete bratu_model;

      if (uses_visit) {
         visit_data_writer.reset();
      }

      t_all->stop();
      int size = tbox::SAMRAI_MPI::getSAMRAIWorld().getSize();
      if (tbox::SAMRAI_MPI::getSAMRAIWorld().getRank() == 0) {
         std::string timing_file =
            base_name + ".timing" + tbox::Utilities::intToString(size);
         FILE* fp = fopen(timing_file.c_str(), "w");
         fprintf(fp, "%f\n", t_all->getTotalWallclockTime());
         fclose(fp);
      }

#endif

#ifdef TESTING
      /*
       * Currently not checking on solution accuracy, so if a run
       * makes it here, it passes.  We need to find a way to check
       * the solution accuracy.
       */
      tbox::pout << "\nPASSED:  nonlinear" << std::endl;
#endif
   }

   tbox::SAMRAIManager::shutdown();
   tbox::SAMRAIManager::finalize();

#ifdef HAVE_PETSC
   PetscFinalize();
#endif
   tbox::SAMRAI_MPI::finalize();

   return 0;
}
